package com.snow.datasource.route;

import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import com.snow.datasource.datasource.DynamicDataSourceContextHolder;
import com.snow.datasource.exception.SnowDataSourceException;
import com.snow.datasource.provider.ISnowProvider;
import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.InitializingBean;

import javax.sql.DataSource;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;

/**
 * @Description
 * @Author fanxl 315620726@qq.com
 * @Date 2020/12/17
 **/
@Data
@Slf4j
public class SnowRoutingDataSource extends AbstractRoutingDataSource implements InitializingBean {

    /**
     * 主数据源
     */
    private String primary;
    /**
     * 数据源
     */
    private Map<String, DataSource> dataSourceMap = new ConcurrentHashMap<>();

    /**
     * 从节点
     */
    private List<String> dataSourceSlaves = new CopyOnWriteArrayList<>();
    /**
     * 数据源属性提供者
     */
    private ISnowProvider iSnowProvider;


    @Override
    protected DataSource determineDataSource() {
        String dataSource = DynamicDataSourceContextHolder.peek();
        if (ObjectUtil.isNull(dataSource)) {
            return getDataSource(primary);
        }
        return getDataSource(DynamicDataSourceContextHolder.peek());
    }

    /**
     * 根据注解获取数据源
     *
     * @param ds
     * @return
     */
    public DataSource getDataSource(String ds) {
        if (StrUtil.isEmpty(ds)) {
            return determineDataSource();
        }
        DataSource dataSource = dataSourceMap.get(ds);
        if (ObjectUtil.isNull(dataSource)) {
            throw new SnowDataSourceException("不存在数据源：" + ds);
        }
        return dataSource;
    }

    /**
     * 删除数据源
     *
     * @param key
     */
    public void removeDataSource(String key) {
        dataSourceMap.remove(key);
    }

    @Override
    public void afterPropertiesSet() {
        this.dataSourceMap = iSnowProvider.loadDataSources();
        for (String key : this.dataSourceMap.keySet()) {
            if (!key.equals(primary)) {
                this.dataSourceSlaves.add(key);
            }
        }

    }
}
